{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c8744d20",
   "metadata": {},
   "source": [
    "### LLM Inference Example\n",
    "\n",
    "This notebook contains a basic inference example for using our `ttml` Python API to build, load, and run a large language model from Hugging Face on our TT hardware. By default, it is set to create and load a GPT2 model, but this notebook can quickly and easily be edited to use any of the LLMs that the tt-train project currently supports. \n",
    "\n",
    "Below, in the first cell, we have our imports and basic directory housekeeping."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "a7a20c4d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os, sys, random\n",
    "import numpy as np  # For numpy arrays\n",
    "from dataclasses import dataclass # For configuration classes\n",
    "from huggingface_hub import hf_hub_download # To download safetensors from Hugging Face\n",
    "from transformers import AutoTokenizer\n",
    "from yaml import safe_load # To read YAML configs\n",
    "from pathlib import Path\n",
    "\n",
    "sys.path.append(f\"{os.environ['TT_METAL_HOME']}/tt-train/sources/ttml\")\n",
    "import ttml\n",
    "from ttml.common.config import get_training_config, load_config\n",
    "from ttml.common.utils import set_seed, round_up_to_tile\n",
    "from ttml.common.model_factory import TransformerModelFactory\n",
    "\n",
    "# Change working directory to tt-train\n",
    "os.chdir(f\"{os.environ['TT_METAL_HOME']}/tt-train\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e005a45c",
   "metadata": {},
   "source": [
    "Use the cell below to change global parameters in this notebook. \n",
    "\n",
    "`OUTPUT_TOKENS` : the length of the generated text in token (not characters!) \n",
    "\n",
    "`WITH_SAMPLING` : enable or disable output token sampling (only used for PyTorch)\n",
    "\n",
    "`TEMPERATURE`   : sampling temperature; set to 0 to disable sampling in `generate_with_tt()`\n",
    "\n",
    "`SEED`          : randomization seed (for reproducibility)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "a87d6dc7",
   "metadata": {},
   "outputs": [],
   "source": [
    "OUTPUT_TOKENS = 256\n",
    "WITH_SAMPLING = True\n",
    "TEMPERATURE = 0.8\n",
    "SEED = 42\n",
    "CONFIG = \"gpt2_inference.yaml\"\n",
    "\n",
    "set_seed(SEED)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9e9b032f",
   "metadata": {},
   "source": [
    "While the notebook is currently configured for GPT2, you can quickly change the tokenizer you want to use by changing the input to `from_pretrained()` below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "ea71a9df",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the tokenizer from Hugging Face and the transformer config from YAML\n",
    "tokenizer = AutoTokenizer.from_pretrained(\"gpt2\")\n",
    "training_config = get_training_config(CONFIG)\n",
    "model_yaml = load_config(training_config.model_config, configs_root=os.getcwd())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5415d24d",
   "metadata": {},
   "source": [
    "As above, the call to `hf_hub_download()` will download (or otherwise find on your local system) the SafeTensors model weight file for GPT2, but can be updated to download other SafeTensors files."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "6965256c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Safetensors path: /home/ubuntu/.cache/huggingface/hub/models--gpt2/snapshots/607a30d783dfa663caf39e06633721c8d4cfcd7e/\n"
     ]
    }
   ],
   "source": [
    "# # Get safetensors\n",
    "safetensors_path = hf_hub_download(repo_id=\"gpt2\", filename=\"model.safetensors\")\n",
    "safetensors_path = safetensors_path.replace(\"model.safetensors\",\"\")\n",
    "\n",
    "print(f\"Safetensors path: {safetensors_path}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "907e5933",
   "metadata": {},
   "outputs": [],
   "source": [
    "def build_causal_mask(T: int) -> ttml.autograd.Tensor:\n",
    "    # [1,1,T,T] float32 with 1s for allowed positions (i >= j), else 0\\n\",\n",
    "    m = np.tril(np.ones((T, T), dtype=np.float32))\n",
    "    return ttml.autograd.Tensor.from_numpy(m.reshape(1, 1, T, T), ttml.Layout.TILE, ttml.autograd.DataType.BFLOAT16)\n",
    "\n",
    "def build_logits_mask(vocab_size: int, padded_vocab_size: int) -> ttml.autograd.Tensor:\n",
    "    logits_mask = np.zeros((1, 1, 1, padded_vocab_size), dtype=np.float32)\n",
    "    logits_mask[:, :, :, vocab_size:] = 1e4\n",
    "    return ttml.autograd.Tensor.from_numpy(logits_mask, ttml.Layout.TILE, ttml.autograd.DataType.BFLOAT16)   # [1,1,1,T], bfloat16\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "bd89028f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Transformer configuration:\n",
      "    Vocab size: 50272\n",
      "    Max sequence length: 1024\n",
      "    Embedding dim: 768\n",
      "    Num heads: 12\n",
      "    Dropout probability: 0.2\n",
      "    Num blocks: 12\n",
      "    Positional embedding type: Trainable\n",
      "    Runner type: Memory efficient\n",
      "    Composite layernorm: false\n",
      "    Weight tying: Enabled\n",
      "2025-12-02 19:15:07.274 | info     |             UMD | Starting topology discovery. (topology_discovery.cpp:69)\n",
      "2025-12-02 19:15:07.278 | info     |             UMD | Established firmware bundle version: 18.10.0 (topology_discovery.cpp:369)\n",
      "2025-12-02 19:15:07.278 | info     |             UMD | Established ETH FW version: 7.0.0 (topology_discovery_wormhole.cpp:324)\n",
      "2025-12-02 19:15:07.278 | info     |             UMD | Completed topology discovery. (topology_discovery.cpp:73)\n",
      "2025-12-02 19:15:07.278 | info     |          Device | Opening user mode device driver (tt_cluster.cpp:211)\n",
      "2025-12-02 19:15:07.278 | info     |             UMD | Starting topology discovery. (topology_discovery.cpp:69)\n",
      "2025-12-02 19:15:07.281 | info     |             UMD | Established firmware bundle version: 18.10.0 (topology_discovery.cpp:369)\n",
      "2025-12-02 19:15:07.281 | info     |             UMD | Established ETH FW version: 7.0.0 (topology_discovery_wormhole.cpp:324)\n",
      "2025-12-02 19:15:07.281 | info     |       Padding vocab size for tilization: original 50257 -> padded 50272\n",
      "      UMD | Completed topology discovery. (topology_discovery.cpp:73)\n",
      "2025-12-02 19:15:07.282 | info     |             UMD | Starting topology discovery. (topology_discovery.cpp:69)\n",
      "2025-12-02 19:15:07.302 | info     |             UMD | Established firmware bundle version: 18.10.0 (topology_discovery.cpp:369)\n",
      "2025-12-02 19:15:07.302 | info     |             UMD | Established ETH FW version: 7.0.0 (topology_discovery_wormhole.cpp:324)\n",
      "2025-12-02 19:15:07.302 | info     |             UMD | Completed topology discovery. (topology_discovery.cpp:73)\n",
      "2025-12-02 19:15:07.302 | info     |             UMD | Harvesting masks for chip 0 tensix: 0x80 dram: 0x0 eth: 0x0 pcie: 0x0 l2cpu: 0x0 (cluster.cpp:358)\n",
      "2025-12-02 19:15:07.371 | warning  |             UMD | init_detect_tt_device_numanodes(): Could not determine NumaNodeSet for TT device (physical_device_id: 0 pci_bus_id: 0000:04:00.0) (cpuset_lib.cpp:564)\n",
      "2025-12-02 19:15:07.371 | warning  |             UMD | Could not find NumaNodeSet for TT Device (physical_device_id: 0 pci_bus_id: 0000:04:00.0) (cpuset_lib.cpp:173)\n",
      "2025-12-02 19:15:07.371 | warning  |             UMD | bind_area_memory_nodeset(): Unable to determine TT Device to NumaNode mapping for physical_device_id: 0. Skipping membind. (cpuset_lib.cpp:450)\n",
      "2025-12-02 19:15:07.371 | warning  |             UMD | ---- ttSiliconDevice::init_hugepage: bind_area_to_memory_nodeset() failed (physical_device_id: 0 ch: 0). Hugepage allocation is not on NumaNode matching TT Device. Side-Effect is decreased Device->Host perf (Issue #893). (sysmem_manager.cpp:219)\n",
      "2025-12-02 19:15:07.372 | info     |             UMD | Opening local chip ids/PCIe ids: {0}/[0] and remote chip ids {} (cluster.cpp:202)\n",
      "2025-12-02 19:15:07.372 | info     |             UMD | IOMMU: disabled (cluster.cpp:178)\n",
      "2025-12-02 19:15:07.372 | info     |             UMD | KMD version: 2.4.0 (cluster.cpp:181)\n",
      "2025-12-02 19:15:07.375 | info     |             UMD | Mapped hugepage 0x200000000 to NOC address 0x800000000 (sysmem_manager.cpp:247)\n",
      "2025-12-02 19:15:07.398 | info     |          Fabric | TopologyMapper mapping start (mesh=0): n_log=1, n_phys=1, log_deg_hist={0:1}, phys_deg_hist={0:1} (topology_mapper_utils.cpp:167)\n",
      "2025-12-02 19:15:07.435 | info     |           Metal | DPRINT enabled on device 0, worker core (x=0,y=0) (virtual (x=18,y=18)). (dprint_server.cpp:689)\n",
      "2025-12-02 19:15:07.436 | info     |           Metal | DPRINT Server attached device 0 (dprint_server.cpp:736)\n",
      "Loading model from: /home/ubuntu/.cache/huggingface/hub/models--gpt2/snapshots/607a30d783dfa663caf39e06633721c8d4cfcd7e/model.safetensors\n",
      "parameter name: transformer/gpt_block_9/mlp/fc2/bias\n",
      "parameter name: transformer/gpt_block_9/mlp/fc1/weight\n",
      "parameter name: transformer/gpt_block_9/attention/qkv_linear/bias\n",
      "parameter name: transformer/gpt_block_9/attention/out_linear/bias\n",
      "parameter name: transformer/gpt_block_8/mlp/fc2/weight\n",
      "parameter name: transformer/gpt_block_8/mlp/fc1/weight\n",
      "parameter name: transformer/gpt_block_8/attention/qkv_linear/bias\n",
      "parameter name: transformer/gpt_block_7/mlp/fc2/weight\n",
      "parameter name: transformer/gpt_block_9/attention/out_linear/weight\n",
      "parameter name: transformer/gpt_block_7/mlp/fc2/bias\n",
      "parameter name: transformer/gpt_block_7/mlp/fc1/weight\n",
      "parameter name: transformer/gpt_block_7/attention/qkv_linear/weight\n",
      "parameter name: transformer/gpt_block_0/ln1/gamma\n",
      "parameter name: transformer/gpt_block_6/attention/out_linear/bias\n",
      "parameter name: transformer/gpt_block_0/attention/qkv_linear/bias\n",
      "parameter name: transformer/gpt_block_4/ln1/beta\n",
      "parameter name: transformer/gpt_block_8/mlp/fc2/bias\n",
      "parameter name: transformer/gpt_block_0/attention/out_linear/weight\n",
      "parameter name: transformer/gpt_block_8/ln2/beta\n",
      "parameter name: transformer/gpt_block_3/attention/out_linear/bias\n",
      "parameter name: transformer/gpt_block_4/mlp/fc2/bias\n",
      "parameter name: transformer/gpt_block_9/ln1/beta\n",
      "parameter name: transformer/gpt_block_5/ln1/gamma\n",
      "parameter name: transformer/gpt_block_11/ln2/beta\n",
      "parameter name: transformer/gpt_block_3/mlp/fc2/bias\n",
      "parameter name: transformer/gpt_block_7/ln2/beta\n",
      "parameter name: transformer/gpt_block_11/attention/qkv_linear/weight\n",
      "parameter name: transformer/gpt_block_0/mlp/fc1/weight\n",
      "parameter name: transformer/gpt_block_7/ln1/beta\n",
      "parameter name: transformer/gpt_block_1/mlp/fc1/weight\n",
      "parameter name: transformer/gpt_block_2/mlp/fc2/bias\n",
      "parameter name: transformer/gpt_block_6/ln1/gamma\n",
      "parameter name: transformer/gpt_block_5/attention/qkv_linear/weight\n",
      "parameter name: transformer/gpt_block_7/mlp/fc1/bias\n",
      "parameter name: transformer/gpt_block_6/ln1/beta\n",
      "parameter name: transformer/gpt_block_5/ln2/gamma\n",
      "parameter name: transformer/gpt_block_0/attention/qkv_linear/weight\n",
      "parameter name: transformer/gpt_block_4/ln2/gamma\n",
      "parameter name: transformer/gpt_block_0/ln2/beta\n",
      "parameter name: transformer/gpt_block_8/ln1/gamma\n",
      "parameter name: transformer/gpt_block_6/mlp/fc1/bias\n",
      "parameter name: transformer/gpt_block_4/ln1/gamma\n",
      "parameter name: transformer/gpt_block_5/ln1/beta\n",
      "parameter name: transformer/gpt_block_7/ln2/gamma\n",
      "parameter name: transformer/gpt_block_3/attention/qkv_linear/bias\n",
      "parameter name: transformer/gpt_block_10/mlp/fc2/weight\n",
      "parameter name: transformer/gpt_block_0/mlp/fc1/bias\n",
      "parameter name: transformer/ln_fc/beta\n",
      "parameter name: transformer/gpt_block_1/ln1/beta\n",
      "parameter name: transformer/gpt_block_4/ln2/beta\n",
      "parameter name: transformer/gpt_block_2/ln2/beta\n",
      "parameter name: transformer/gpt_block_1/ln2/beta\n",
      "parameter name: transformer/gpt_block_7/ln1/gamma\n",
      "parameter name: transformer/gpt_block_3/ln1/gamma\n",
      "parameter name: transformer/gpt_block_8/ln2/gamma\n",
      "parameter name: transformer/gpt_block_11/mlp/fc2/bias\n",
      "parameter name: transformer/gpt_block_1/ln2/gamma\n",
      "parameter name: transformer/gpt_block_0/attention/out_linear/bias\n",
      "parameter name: transformer/gpt_block_2/attention/qkv_linear/bias\n",
      "parameter name: transformer/gpt_block_10/ln2/gamma\n",
      "parameter name: transformer/pos_emb/weight\n",
      "parameter name: transformer/gpt_block_5/mlp/fc2/bias\n",
      "parameter name: transformer/gpt_block_11/ln1/beta\n",
      "parameter name: transformer/gpt_block_9/attention/qkv_linear/weight\n",
      "parameter name: transformer/gpt_block_0/ln2/gamma\n",
      "parameter name: transformer/gpt_block_11/ln1/gamma\n",
      "parameter name: transformer/gpt_block_1/attention/out_linear/weight\n",
      "parameter name: transformer/gpt_block_5/mlp/fc2/weight\n",
      "parameter name: transformer/gpt_block_9/ln2/gamma\n",
      "parameter name: transformer/gpt_block_10/ln1/gamma\n",
      "parameter name: transformer/gpt_block_4/attention/out_linear/weight\n",
      "parameter name: transformer/gpt_block_9/ln1/gamma\n",
      "parameter name: transformer/gpt_block_0/mlp/fc2/weight\n",
      "parameter name: transformer/gpt_block_5/attention/out_linear/weight\n",
      "parameter name: transformer/gpt_block_10/mlp/fc1/bias\n",
      "parameter name: transformer/gpt_block_11/mlp/fc1/bias\n",
      "parameter name: transformer/gpt_block_4/mlp/fc1/weight\n",
      "parameter name: transformer/gpt_block_2/ln1/gamma\n",
      "parameter name: transformer/fc/weight\n",
      "parameter name: transformer/gpt_block_11/ln2/gamma\n",
      "parameter name: transformer/ln_fc/gamma\n",
      "parameter name: transformer/gpt_block_8/attention/out_linear/weight\n",
      "parameter name: transformer/gpt_block_9/ln2/beta\n",
      "parameter name: transformer/gpt_block_6/attention/out_linear/weight\n",
      "parameter name: transformer/gpt_block_6/ln2/gamma\n",
      "parameter name: transformer/gpt_block_4/attention/qkv_linear/bias\n",
      "parameter name: transformer/gpt_block_10/ln2/beta\n",
      "parameter name: transformer/gpt_block_6/mlp/fc1/weight\n",
      "parameter name: transformer/gpt_block_2/ln2/gamma\n",
      "parameter name: transformer/gpt_block_1/ln1/gamma\n",
      "parameter name: transformer/gpt_block_0/mlp/fc2/bias\n",
      "parameter name: transformer/gpt_block_1/mlp/fc2/bias\n",
      "parameter name: transformer/gpt_block_3/ln2/gamma\n",
      "parameter name: transformer/gpt_block_9/mlp/fc1/bias\n",
      "parameter name: transformer/gpt_block_3/ln1/beta\n",
      "parameter name: transformer/gpt_block_1/attention/out_linear/bias\n",
      "parameter name: transformer/gpt_block_2/mlp/fc2/weight\n",
      "parameter name: transformer/gpt_block_1/attention/qkv_linear/weight\n",
      "parameter name: transformer/gpt_block_1/mlp/fc1/bias\n",
      "parameter name: transformer/gpt_block_3/attention/out_linear/weight\n",
      "parameter name: transformer/gpt_block_4/attention/qkv_linear/weight\n",
      "parameter name: transformer/gpt_block_8/attention/out_linear/bias\n",
      "parameter name: transformer/gpt_block_11/mlp/fc1/weight\n",
      "parameter name: transformer/gpt_block_1/mlp/fc2/weight\n",
      "parameter name: transformer/gpt_block_2/mlp/fc1/weight\n",
      "parameter name: transformer/gpt_block_1/attention/qkv_linear/bias\n",
      "parameter name: transformer/gpt_block_10/mlp/fc1/weight\n",
      "parameter name: transformer/gpt_block_10/attention/out_linear/bias\n",
      "parameter name: transformer/gpt_block_11/attention/out_linear/bias\n",
      "parameter name: transformer/gpt_block_8/attention/qkv_linear/weight\n",
      "parameter name: transformer/gpt_block_2/ln1/beta\n",
      "parameter name: transformer/gpt_block_10/attention/out_linear/weight\n",
      "parameter name: transformer/gpt_block_5/ln2/beta\n",
      "parameter name: transformer/gpt_block_11/mlp/fc2/weight\n",
      "parameter name: transformer/gpt_block_10/attention/qkv_linear/bias\n",
      "parameter name: transformer/gpt_block_10/mlp/fc2/bias\n",
      "parameter name: transformer/gpt_block_4/mlp/fc1/bias\n",
      "parameter name: transformer/gpt_block_11/attention/out_linear/weight\n",
      "parameter name: transformer/gpt_block_3/mlp/fc2/weight\n",
      "parameter name: transformer/gpt_block_11/attention/qkv_linear/bias\n",
      "parameter name: transformer/gpt_block_5/mlp/fc1/weight\n",
      "parameter name: transformer/gpt_block_2/attention/out_linear/bias\n",
      "parameter name: transformer/gpt_block_2/attention/out_linear/weight\n",
      "parameter name: transformer/gpt_block_6/ln2/beta\n",
      "parameter name: transformer/gpt_block_2/attention/qkv_linear/weight\n",
      "parameter name: transformer/gpt_block_0/ln1/beta\n",
      "parameter name: transformer/gpt_block_10/attention/qkv_linear/weight\n",
      "parameter name: transformer/gpt_block_2/mlp/fc1/bias\n",
      "parameter name: transformer/gpt_block_3/attention/qkv_linear/weight\n",
      "parameter name: transformer/gpt_block_3/mlp/fc1/bias\n",
      "parameter name: transformer/gpt_block_8/ln1/beta\n",
      "parameter name: transformer/gpt_block_3/mlp/fc1/weight\n",
      "parameter name: transformer/gpt_block_6/attention/qkv_linear/bias\n",
      "parameter name: transformer/gpt_block_4/attention/out_linear/bias\n",
      "parameter name: transformer/gpt_block_4/mlp/fc2/weight\n",
      "parameter name: transformer/gpt_block_9/mlp/fc2/weight\n",
      "parameter name: transformer/gpt_block_10/ln1/beta\n",
      "parameter name: transformer/gpt_block_5/attention/out_linear/bias\n",
      "parameter name: transformer/gpt_block_5/attention/qkv_linear/bias\n",
      "parameter name: transformer/gpt_block_5/mlp/fc1/bias\n",
      "parameter name: transformer/gpt_block_6/attention/qkv_linear/weight\n",
      "parameter name: transformer/gpt_block_6/mlp/fc2/bias\n",
      "parameter name: transformer/gpt_block_8/mlp/fc1/bias\n",
      "parameter name: transformer/gpt_block_7/attention/qkv_linear/bias\n",
      "parameter name: transformer/gpt_block_3/ln2/beta\n",
      "parameter name: transformer/gpt_block_6/mlp/fc2/weight\n",
      "parameter name: transformer/gpt_block_7/attention/out_linear/bias\n",
      "parameter name: transformer/gpt_block_7/attention/out_linear/weight\n",
      "Loading tensor: h.0.attn.bias, shape:Shape([1, 1, 1024, 1024]), format: F32\n",
      "Loading tensor: h.0.attn.c_attn.bias, shape:Shape([2304]), format: F32\n",
      " Parameter transformer/gpt_block_0/attention/qkv_linear/bias, shape: Shape([1, 1, 1, 2304])\n",
      "Loading tensor: h.0.attn.c_attn.weight, shape:Shape([768, 2304]), format: F32\n",
      "Transposing!\n",
      " Parameter transformer/gpt_block_0/attention/qkv_linear/weight, shape: Shape([1, 1, 2304, 768])\n",
      "Loading tensor: h.0.attn.c_proj.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_0/attention/out_linear/bias, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.0.attn.c_proj.weight, shape:Shape([768, 768]), format: F32\n",
      "Transposing!\n",
      " Parameter transformer/gpt_block_0/attention/out_linear/weight, shape: Shape([1, 1, 768, 768])\n",
      "Loading tensor: h.0.ln_1.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_0/ln1/beta, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.0.ln_1.weight, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_0/ln1/gamma, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.0.ln_2.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_0/ln2/beta, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.0.ln_2.weight, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_0/ln2/gamma, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.0.mlp.c_fc.bias, shape:Shape([3072]), format: F32\n",
      " Parameter transformer/gpt_block_0/mlp/fc1/bias, shape: Shape([1, 1, 1, 3072])\n",
      "Loading tensor: h.0.mlp.c_fc.weight, shape:Shape([768, 3072]), format: F32\n",
      " Parameter transformer/gpt_block_0/mlp/fc1/weight, shape: Shape([1, 1, 3072, 768])\n",
      "Transposing!\n",
      "Loading tensor: h.0.mlp.c_proj.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_0/mlp/fc2/bias, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.0.mlp.c_proj.weight, shape:Shape([3072, 768]), format: F32\n",
      " Parameter transformer/gpt_block_0/mlp/fc2/weight, shape: Shape([1, 1, 768, 3072])\n",
      "Transposing!\n",
      "Loading tensor: h.1.attn.bias, shape:Shape([1, 1, 1024, 1024]), format: F32\n",
      "Loading tensor: h.1.attn.c_attn.bias, shape:Shape([2304]), format: F32\n",
      " Parameter transformer/gpt_block_1/attention/qkv_linear/bias, shape: Shape([1, 1, 1, 2304])\n",
      "Loading tensor: h.1.attn.c_attn.weight, shape:Shape([768, 2304]), format: F32\n",
      "Transposing!\n",
      " Parameter transformer/gpt_block_1/attention/qkv_linear/weight, shape: Shape([1, 1, 2304, 768])\n",
      "Loading tensor: h.1.attn.c_proj.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_1/attention/out_linear/bias, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.1.attn.c_proj.weight, shape:Shape([768, 768]), format: F32\n",
      "Transposing!\n",
      " Parameter transformer/gpt_block_1/attention/out_linear/weight, shape: Shape([1, 1, 768, 768])\n",
      "Loading tensor: h.1.ln_1.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_1/ln1/beta, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.1.ln_1.weight, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_1/ln1/gamma, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.1.ln_2.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_1/ln2/beta, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.1.ln_2.weight, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_1/ln2/gamma, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.1.mlp.c_fc.bias, shape:Shape([3072]), format: F32\n",
      " Parameter transformer/gpt_block_1/mlp/fc1/bias, shape: Shape([1, 1, 1, 3072])\n",
      "Loading tensor: h.1.mlp.c_fc.weight, shape:Shape([768, 3072]), format: F32\n",
      " Parameter transformer/gpt_block_1/mlp/fc1/weight, shape: Shape([1, 1, 3072, 768])\n",
      "Transposing!\n",
      "Loading tensor: h.1.mlp.c_proj.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_1/mlp/fc2/bias, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.1.mlp.c_proj.weight, shape:Shape([3072, 768]), format: F32\n",
      " Parameter transformer/gpt_block_1/mlp/fc2/weight, shape: Shape([1, 1, 768, 3072])\n",
      "Transposing!\n",
      "Loading tensor: h.10.attn.bias, shape:Shape([1, 1, 1024, 1024]), format: F32\n",
      "Loading tensor: h.10.attn.c_attn.bias, shape:Shape([2304]), format: F32\n",
      " Parameter transformer/gpt_block_10/attention/qkv_linear/bias, shape: Shape([1, 1, 1, 2304])\n",
      "Loading tensor: h.10.attn.c_attn.weight, shape:Shape([768, 2304]), format: F32\n",
      "Transposing!\n",
      " Parameter transformer/gpt_block_10/attention/qkv_linear/weight, shape: Shape([1, 1, 2304, 768])\n",
      "Loading tensor: h.10.attn.c_proj.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_10/attention/out_linear/bias, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.10.attn.c_proj.weight, shape:Shape([768, 768]), format: F32\n",
      "Transposing!\n",
      " Parameter transformer/gpt_block_10/attention/out_linear/weight, shape: Shape([1, 1, 768, 768])\n",
      "Loading tensor: h.10.ln_1.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_10/ln1/beta, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.10.ln_1.weight, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_10/ln1/gamma, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.10.ln_2.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_10/ln2/beta, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.10.ln_2.weight, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_10/ln2/gamma, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.10.mlp.c_fc.bias, shape:Shape([3072]), format: F32\n",
      " Parameter transformer/gpt_block_10/mlp/fc1/bias, shape: Shape([1, 1, 1, 3072])\n",
      "Loading tensor: h.10.mlp.c_fc.weight, shape:Shape([768, 3072]), format: F32\n",
      " Parameter transformer/gpt_block_10/mlp/fc1/weight, shape: Shape([1, 1, 3072, 768])\n",
      "Transposing!\n",
      "Loading tensor: h.10.mlp.c_proj.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_10/mlp/fc2/bias, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.10.mlp.c_proj.weight, shape:Shape([3072, 768]), format: F32\n",
      " Parameter transformer/gpt_block_10/mlp/fc2/weight, shape: Shape([1, 1, 768, 3072])\n",
      "Transposing!\n",
      "Loading tensor: h.11.attn.bias, shape:Shape([1, 1, 1024, 1024]), format: F32\n",
      "Loading tensor: h.11.attn.c_attn.bias, shape:Shape([2304]), format: F32\n",
      " Parameter transformer/gpt_block_11/attention/qkv_linear/bias, shape: Shape([1, 1, 1, 2304])\n",
      "Loading tensor: h.11.attn.c_attn.weight, shape:Shape([768, 2304]), format: F32\n",
      "Transposing!\n",
      " Parameter transformer/gpt_block_11/attention/qkv_linear/weight, shape: Shape([1, 1, 2304, 768])\n",
      "Loading tensor: h.11.attn.c_proj.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_11/attention/out_linear/bias, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.11.attn.c_proj.weight, shape:Shape([768, 768]), format: F32\n",
      "Transposing!\n",
      " Parameter transformer/gpt_block_11/attention/out_linear/weight, shape: Shape([1, 1, 768, 768])\n",
      "Loading tensor: h.11.ln_1.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_11/ln1/beta, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.11.ln_1.weight, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_11/ln1/gamma, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.11.ln_2.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_11/ln2/beta, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.11.ln_2.weight, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_11/ln2/gamma, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.11.mlp.c_fc.bias, shape:Shape([3072]), format: F32\n",
      " Parameter transformer/gpt_block_11/mlp/fc1/bias, shape: Shape([1, 1, 1, 3072])\n",
      "Loading tensor: h.11.mlp.c_fc.weight, shape:Shape([768, 3072]), format: F32\n",
      " Parameter transformer/gpt_block_11/mlp/fc1/weight, shape: Shape([1, 1, 3072, 768])\n",
      "Transposing!\n",
      "Loading tensor: h.11.mlp.c_proj.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_11/mlp/fc2/bias, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.11.mlp.c_proj.weight, shape:Shape([3072, 768]), format: F32\n",
      " Parameter transformer/gpt_block_11/mlp/fc2/weight, shape: Shape([1, 1, 768, 3072])\n",
      "Transposing!\n",
      "Loading tensor: h.2.attn.bias, shape:Shape([1, 1, 1024, 1024]), format: F32\n",
      "Loading tensor: h.2.attn.c_attn.bias, shape:Shape([2304]), format: F32\n",
      " Parameter transformer/gpt_block_2/attention/qkv_linear/bias, shape: Shape([1, 1, 1, 2304])\n",
      "Loading tensor: h.2.attn.c_attn.weight, shape:Shape([768, 2304]), format: F32\n",
      "Transposing!\n",
      " Parameter transformer/gpt_block_2/attention/qkv_linear/weight, shape: Shape([1, 1, 2304, 768])\n",
      "Loading tensor: h.2.attn.c_proj.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_2/attention/out_linear/bias, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.2.attn.c_proj.weight, shape:Shape([768, 768]), format: F32\n",
      "Transposing!\n",
      " Parameter transformer/gpt_block_2/attention/out_linear/weight, shape: Shape([1, 1, 768, 768])\n",
      "Loading tensor: h.2.ln_1.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_2/ln1/beta, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.2.ln_1.weight, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_2/ln1/gamma, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.2.ln_2.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_2/ln2/beta, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.2.ln_2.weight, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_2/ln2/gamma, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.2.mlp.c_fc.bias, shape:Shape([3072]), format: F32\n",
      " Parameter transformer/gpt_block_2/mlp/fc1/bias, shape: Shape([1, 1, 1, 3072])\n",
      "Loading tensor: h.2.mlp.c_fc.weight, shape:Shape([768, 3072]), format: F32\n",
      " Parameter transformer/gpt_block_2/mlp/fc1/weight, shape: Shape([1, 1, 3072, 768])\n",
      "Transposing!\n",
      "Loading tensor: h.2.mlp.c_proj.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_2/mlp/fc2/bias, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.2.mlp.c_proj.weight, shape:Shape([3072, 768]), format: F32\n",
      " Parameter transformer/gpt_block_2/mlp/fc2/weight, shape: Shape([1, 1, 768, 3072])\n",
      "Transposing!\n",
      "Loading tensor: h.3.attn.bias, shape:Shape([1, 1, 1024, 1024]), format: F32\n",
      "Loading tensor: h.3.attn.c_attn.bias, shape:Shape([2304]), format: F32\n",
      " Parameter transformer/gpt_block_3/attention/qkv_linear/bias, shape: Shape([1, 1, 1, 2304])\n",
      "Loading tensor: h.3.attn.c_attn.weight, shape:Shape([768, 2304]), format: F32\n",
      "Transposing!\n",
      " Parameter transformer/gpt_block_3/attention/qkv_linear/weight, shape: Shape([1, 1, 2304, 768])\n",
      "Loading tensor: h.3.attn.c_proj.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_3/attention/out_linear/bias, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.3.attn.c_proj.weight, shape:Shape([768, 768]), format: F32\n",
      "Transposing!\n",
      " Parameter transformer/gpt_block_3/attention/out_linear/weight, shape: Shape([1, 1, 768, 768])\n",
      "Loading tensor: h.3.ln_1.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_3/ln1/beta, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.3.ln_1.weight, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_3/ln1/gamma, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.3.ln_2.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_3/ln2/beta, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.3.ln_2.weight, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_3/ln2/gamma, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.3.mlp.c_fc.bias, shape:Shape([3072]), format: F32\n",
      " Parameter transformer/gpt_block_3/mlp/fc1/bias, shape: Shape([1, 1, 1, 3072])\n",
      "Loading tensor: h.3.mlp.c_fc.weight, shape:Shape([768, 3072]), format: F32\n",
      " Parameter transformer/gpt_block_3/mlp/fc1/weight, shape: Shape([1, 1, 3072, 768])\n",
      "Transposing!\n",
      "Loading tensor: h.3.mlp.c_proj.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_3/mlp/fc2/bias, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.3.mlp.c_proj.weight, shape:Shape([3072, 768]), format: F32\n",
      " Parameter transformer/gpt_block_3/mlp/fc2/weight, shape: Shape([1, 1, 768, 3072])\n",
      "Transposing!\n",
      "Loading tensor: h.4.attn.bias, shape:Shape([1, 1, 1024, 1024]), format: F32\n",
      "Loading tensor: h.4.attn.c_attn.bias, shape:Shape([2304]), format: F32\n",
      " Parameter transformer/gpt_block_4/attention/qkv_linear/bias, shape: Shape([1, 1, 1, 2304])\n",
      "Loading tensor: h.4.attn.c_attn.weight, shape:Shape([768, 2304]), format: F32\n",
      "Transposing!\n",
      " Parameter transformer/gpt_block_4/attention/qkv_linear/weight, shape: Shape([1, 1, 2304, 768])\n",
      "Loading tensor: h.4.attn.c_proj.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_4/attention/out_linear/bias, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.4.attn.c_proj.weight, shape:Shape([768, 768]), format: F32\n",
      "Transposing!\n",
      " Parameter transformer/gpt_block_4/attention/out_linear/weight, shape: Shape([1, 1, 768, 768])\n",
      "Loading tensor: h.4.ln_1.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_4/ln1/beta, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.4.ln_1.weight, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_4/ln1/gamma, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.4.ln_2.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_4/ln2/beta, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.4.ln_2.weight, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_4/ln2/gamma, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.4.mlp.c_fc.bias, shape:Shape([3072]), format: F32\n",
      " Parameter transformer/gpt_block_4/mlp/fc1/bias, shape: Shape([1, 1, 1, 3072])\n",
      "Loading tensor: h.4.mlp.c_fc.weight, shape:Shape([768, 3072]), format: F32\n",
      " Parameter transformer/gpt_block_4/mlp/fc1/weight, shape: Shape([1, 1, 3072, 768])\n",
      "Transposing!\n",
      "Loading tensor: h.4.mlp.c_proj.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_4/mlp/fc2/bias, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.4.mlp.c_proj.weight, shape:Shape([3072, 768]), format: F32\n",
      " Parameter transformer/gpt_block_4/mlp/fc2/weight, shape: Shape([1, 1, 768, 3072])\n",
      "Transposing!\n",
      "Loading tensor: h.5.attn.bias, shape:Shape([1, 1, 1024, 1024]), format: F32\n",
      "Loading tensor: h.5.attn.c_attn.bias, shape:Shape([2304]), format: F32\n",
      " Parameter transformer/gpt_block_5/attention/qkv_linear/bias, shape: Shape([1, 1, 1, 2304])\n",
      "Loading tensor: h.5.attn.c_attn.weight, shape:Shape([768, 2304]), format: F32\n",
      "Transposing!\n",
      " Parameter transformer/gpt_block_5/attention/qkv_linear/weight, shape: Shape([1, 1, 2304, 768])\n",
      "Loading tensor: h.5.attn.c_proj.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_5/attention/out_linear/bias, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.5.attn.c_proj.weight, shape:Shape([768, 768]), format: F32\n",
      "Transposing!\n",
      " Parameter transformer/gpt_block_5/attention/out_linear/weight, shape: Shape([1, 1, 768, 768])\n",
      "Loading tensor: h.5.ln_1.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_5/ln1/beta, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.5.ln_1.weight, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_5/ln1/gamma, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.5.ln_2.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_5/ln2/beta, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.5.ln_2.weight, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_5/ln2/gamma, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.5.mlp.c_fc.bias, shape:Shape([3072]), format: F32\n",
      " Parameter transformer/gpt_block_5/mlp/fc1/bias, shape: Shape([1, 1, 1, 3072])\n",
      "Loading tensor: h.5.mlp.c_fc.weight, shape:Shape([768, 3072]), format: F32\n",
      " Parameter transformer/gpt_block_5/mlp/fc1/weight, shape: Shape([1, 1, 3072, 768])\n",
      "Transposing!\n",
      "Loading tensor: h.5.mlp.c_proj.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_5/mlp/fc2/bias, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.5.mlp.c_proj.weight, shape:Shape([3072, 768]), format: F32\n",
      " Parameter transformer/gpt_block_5/mlp/fc2/weight, shape: Shape([1, 1, 768, 3072])\n",
      "Transposing!\n",
      "Loading tensor: h.6.attn.bias, shape:Shape([1, 1, 1024, 1024]), format: F32\n",
      "Loading tensor: h.6.attn.c_attn.bias, shape:Shape([2304]), format: F32\n",
      " Parameter transformer/gpt_block_6/attention/qkv_linear/bias, shape: Shape([1, 1, 1, 2304])\n",
      "Loading tensor: h.6.attn.c_attn.weight, shape:Shape([768, 2304]), format: F32\n",
      "Transposing!\n",
      " Parameter transformer/gpt_block_6/attention/qkv_linear/weight, shape: Shape([1, 1, 2304, 768])\n",
      "Loading tensor: h.6.attn.c_proj.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_6/attention/out_linear/bias, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.6.attn.c_proj.weight, shape:Shape([768, 768]), format: F32\n",
      "Transposing!\n",
      " Parameter transformer/gpt_block_6/attention/out_linear/weight, shape: Shape([1, 1, 768, 768])\n",
      "Loading tensor: h.6.ln_1.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_6/ln1/beta, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.6.ln_1.weight, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_6/ln1/gamma, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.6.ln_2.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_6/ln2/beta, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.6.ln_2.weight, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_6/ln2/gamma, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.6.mlp.c_fc.bias, shape:Shape([3072]), format: F32\n",
      " Parameter transformer/gpt_block_6/mlp/fc1/bias, shape: Shape([1, 1, 1, 3072])\n",
      "Loading tensor: h.6.mlp.c_fc.weight, shape:Shape([768, 3072]), format: F32\n",
      " Parameter transformer/gpt_block_6/mlp/fc1/weight, shape: Shape([1, 1, 3072, 768])\n",
      "Transposing!\n",
      "Loading tensor: h.6.mlp.c_proj.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_6/mlp/fc2/bias, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.6.mlp.c_proj.weight, shape:Shape([3072, 768]), format: F32\n",
      " Parameter transformer/gpt_block_6/mlp/fc2/weight, shape: Shape([1, 1, 768, 3072])\n",
      "Transposing!\n",
      "Loading tensor: h.7.attn.bias, shape:Shape([1, 1, 1024, 1024]), format: F32\n",
      "Loading tensor: h.7.attn.c_attn.bias, shape:Shape([2304]), format: F32\n",
      " Parameter transformer/gpt_block_7/attention/qkv_linear/bias, shape: Shape([1, 1, 1, 2304])\n",
      "Loading tensor: h.7.attn.c_attn.weight, shape:Shape([768, 2304]), format: F32\n",
      "Transposing!\n",
      " Parameter transformer/gpt_block_7/attention/qkv_linear/weight, shape: Shape([1, 1, 2304, 768])\n",
      "Loading tensor: h.7.attn.c_proj.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_7/attention/out_linear/bias, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.7.attn.c_proj.weight, shape:Shape([768, 768]), format: F32\n",
      "Transposing!\n",
      " Parameter transformer/gpt_block_7/attention/out_linear/weight, shape: Shape([1, 1, 768, 768])\n",
      "Loading tensor: h.7.ln_1.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_7/ln1/beta, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.7.ln_1.weight, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_7/ln1/gamma, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.7.ln_2.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_7/ln2/beta, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.7.ln_2.weight, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_7/ln2/gamma, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.7.mlp.c_fc.bias, shape:Shape([3072]), format: F32\n",
      " Parameter transformer/gpt_block_7/mlp/fc1/bias, shape: Shape([1, 1, 1, 3072])\n",
      "Loading tensor: h.7.mlp.c_fc.weight, shape:Shape([768, 3072]), format: F32\n",
      " Parameter transformer/gpt_block_7/mlp/fc1/weight, shape: Shape([1, 1, 3072, 768])\n",
      "Transposing!\n",
      "Loading tensor: h.7.mlp.c_proj.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_7/mlp/fc2/bias, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.7.mlp.c_proj.weight, shape:Shape([3072, 768]), format: F32\n",
      " Parameter transformer/gpt_block_7/mlp/fc2/weight, shape: Shape([1, 1, 768, 3072])\n",
      "Transposing!\n",
      "Loading tensor: h.8.attn.bias, shape:Shape([1, 1, 1024, 1024]), format: F32\n",
      "Loading tensor: h.8.attn.c_attn.bias, shape:Shape([2304]), format: F32\n",
      " Parameter transformer/gpt_block_8/attention/qkv_linear/bias, shape: Shape([1, 1, 1, 2304])\n",
      "Loading tensor: h.8.attn.c_attn.weight, shape:Shape([768, 2304]), format: F32\n",
      "Transposing!\n",
      " Parameter transformer/gpt_block_8/attention/qkv_linear/weight, shape: Shape([1, 1, 2304, 768])\n",
      "Loading tensor: h.8.attn.c_proj.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_8/attention/out_linear/bias, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.8.attn.c_proj.weight, shape:Shape([768, 768]), format: F32\n",
      "Transposing!\n",
      " Parameter transformer/gpt_block_8/attention/out_linear/weight, shape: Shape([1, 1, 768, 768])\n",
      "Loading tensor: h.8.ln_1.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_8/ln1/beta, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.8.ln_1.weight, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_8/ln1/gamma, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.8.ln_2.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_8/ln2/beta, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.8.ln_2.weight, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_8/ln2/gamma, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.8.mlp.c_fc.bias, shape:Shape([3072]), format: F32\n",
      " Parameter transformer/gpt_block_8/mlp/fc1/bias, shape: Shape([1, 1, 1, 3072])\n",
      "Loading tensor: h.8.mlp.c_fc.weight, shape:Shape([768, 3072]), format: F32\n",
      " Parameter transformer/gpt_block_8/mlp/fc1/weight, shape: Shape([1, 1, 3072, 768])\n",
      "Transposing!\n",
      "Loading tensor: h.8.mlp.c_proj.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_8/mlp/fc2/bias, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.8.mlp.c_proj.weight, shape:Shape([3072, 768]), format: F32\n",
      " Parameter transformer/gpt_block_8/mlp/fc2/weight, shape: Shape([1, 1, 768, 3072])\n",
      "Transposing!\n",
      "Loading tensor: h.9.attn.bias, shape:Shape([1, 1, 1024, 1024]), format: F32\n",
      "Loading tensor: h.9.attn.c_attn.bias, shape:Shape([2304]), format: F32\n",
      " Parameter transformer/gpt_block_9/attention/qkv_linear/bias, shape: Shape([1, 1, 1, 2304])\n",
      "Loading tensor: h.9.attn.c_attn.weight, shape:Shape([768, 2304]), format: F32\n",
      "Transposing!\n",
      " Parameter transformer/gpt_block_9/attention/qkv_linear/weight, shape: Shape([1, 1, 2304, 768])\n",
      "Loading tensor: h.9.attn.c_proj.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_9/attention/out_linear/bias, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.9.attn.c_proj.weight, shape:Shape([768, 768]), format: F32\n",
      "Transposing!\n",
      " Parameter transformer/gpt_block_9/attention/out_linear/weight, shape: Shape([1, 1, 768, 768])\n",
      "Loading tensor: h.9.ln_1.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_9/ln1/beta, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.9.ln_1.weight, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_9/ln1/gamma, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.9.ln_2.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_9/ln2/beta, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.9.ln_2.weight, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_9/ln2/gamma, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.9.mlp.c_fc.bias, shape:Shape([3072]), format: F32\n",
      " Parameter transformer/gpt_block_9/mlp/fc1/bias, shape: Shape([1, 1, 1, 3072])\n",
      "Loading tensor: h.9.mlp.c_fc.weight, shape:Shape([768, 3072]), format: F32\n",
      " Parameter transformer/gpt_block_9/mlp/fc1/weight, shape: Shape([1, 1, 3072, 768])\n",
      "Transposing!\n",
      "Loading tensor: h.9.mlp.c_proj.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_9/mlp/fc2/bias, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.9.mlp.c_proj.weight, shape:Shape([3072, 768]), format: F32\n",
      " Parameter transformer/gpt_block_9/mlp/fc2/weight, shape: Shape([1, 1, 768, 3072])\n",
      "Transposing!\n",
      "Loading tensor: ln_f.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/ln_fc/beta, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: ln_f.weight, shape:Shape([768]), format: F32\n",
      " Parameter transformer/ln_fc/gamma, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: wpe.weight, shape:Shape([1024, 768]), format: F32\n",
      " Parameter transformer/pos_emb/weight, shape: Shape([1, 1, 1024, 768])\n",
      "Loading tensor: wte.weight, shape:Shape([50257, 768]), format: F32\n",
      " Parameter transformer/fc/weight, shape: Shape([1, 1, 50272, 768])\n",
      "Original shape 50257, 768Transformed shape 50272, 768\n"
     ]
    }
   ],
   "source": [
    "orig_vocab_size = tokenizer.vocab_size\n",
    "\n",
    "tt_model_factory = TransformerModelFactory(model_yaml)\n",
    "tt_model_factory.transformer_config.vocab_size = orig_vocab_size\n",
    "\n",
    "max_sequence_length = tt_model_factory.transformer_config.max_sequence_length\n",
    "\n",
    "tt_model = tt_model_factory.create_model()\n",
    "tt_model.load_from_safetensors(safetensors_path)\n",
    "tt_model\n",
    "\n",
    "padded_vocab_size = round_up_to_tile(orig_vocab_size, 32)\n",
    "\n",
    "if orig_vocab_size != padded_vocab_size:\n",
    "    print(f\"Padding vocab size for tilization: original {orig_vocab_size} -> padded {padded_vocab_size}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "40913e88",
   "metadata": {},
   "source": [
    "`generate_with_tt()` uses TT hardware acceleration to generate output from the chosen LLM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "18e6550e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_with_tt(model, prompt_tokens):\n",
    "\n",
    "    ttml.autograd.AutoContext.get_instance().set_gradient_mode(ttml.autograd.GradMode.DISABLED)\n",
    "    model.eval()\n",
    "\n",
    "    logits_mask_tensor = None\n",
    "\n",
    "    if padded_vocab_size != orig_vocab_size:\n",
    "        logits_mask_tensor = build_logits_mask(orig_vocab_size, padded_vocab_size)\n",
    "\n",
    "    causal_mask = build_causal_mask(max_sequence_length)  # [1,1,seq_len,seq_len], float32\n",
    "    padded_prompt_tokens = np.zeros((1, 1, 1, max_sequence_length), \n",
    "                                    dtype=np.uint32)\n",
    "\n",
    "    start_idx = 0\n",
    "\n",
    "    print(\"************************************\")\n",
    "    for token_idx in range(OUTPUT_TOKENS):\n",
    "\n",
    "        if len(prompt_tokens) > max_sequence_length:\n",
    "            start_idx = len(prompt_tokens) - max_sequence_length\n",
    "\n",
    "        # padded_prompt_tokens[0, 0, 0, :transformer_cfg[\"max_sequence_length\"]] = 0\n",
    "        padded_prompt_tokens[0, 0, 0, :len(prompt_tokens)] = prompt_tokens[start_idx:]\n",
    "        padded_prompt_tensor = ttml.autograd.Tensor.from_numpy(\n",
    "            padded_prompt_tokens,\n",
    "            ttml.Layout.ROW_MAJOR,\n",
    "            ttml.autograd.DataType.UINT32)  # [1,1,1, max_seq_len], uint32\n",
    "\n",
    "        logits = model(padded_prompt_tensor, causal_mask)  # out=[1,1,seq_len, vocab_size], bf16\n",
    "\n",
    "\n",
    "        next_token_tensor = ttml.ops.sample.sample_op(logits, TEMPERATURE, np.random.randint(low=1e7), logits_mask_tensor)  # out=[1,1,seq_len,1], uint32\n",
    "\n",
    "        next_token_idx = max_sequence_length - 1 if len(prompt_tokens) > max_sequence_length else len(prompt_tokens) - 1\n",
    "        next_token = next_token_tensor.to_numpy().flatten()[next_token_idx]\n",
    "\n",
    "        output = tokenizer.decode(next_token)\n",
    "\n",
    "        prompt_tokens.append(next_token)\n",
    "        print(output, end='', flush=True)\n",
    "\n",
    "    print(\"\\n************************************\\n\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "7824c0cb",
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_with_pytorch(prompt_tokens):\n",
    "    import torch\n",
    "    from transformers import AutoModelForCausalLM\n",
    "\n",
    "    torch.manual_seed(SEED)\n",
    "\n",
    "    torch_model = AutoModelForCausalLM.from_pretrained(\"gpt2\", dtype=torch.bfloat16)\n",
    "    torch_model.eval()\n",
    "    print(\"************************************\")\n",
    "    with torch.no_grad():\n",
    "        outputs = torch_model.generate(\n",
    "            prompt_tokens,\n",
    "            max_new_tokens=OUTPUT_TOKENS,\n",
    "            do_sample=WITH_SAMPLING, # Enable sampling\n",
    "            temperature=TEMPERATURE,   # Temperature for sampling\n",
    "            num_beams=1 # Use multinomial sampling (standard sampling)\n",
    "        )\n",
    "    \n",
    "    generated_text = tokenizer.batch_decode(outputs, skip_special_tokens=True)\n",
    "\n",
    "    for t in generated_text:\n",
    "        print(t)\n",
    "        \n",
    "    print(\"\\n************************************\\n\\n\"),"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6a42273b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Generating with TT:\n"
     ]
    }
   ],
   "source": [
    "prompt_str = \"The difference between cats and dogs is:\"\n",
    "prompt_tokens = tokenizer.encode(prompt_str)\n",
    "print(\"Generating with TT:\")\n",
    "generate_with_tt(tt_model, prompt_tokens.copy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "d98a32d0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Generating with TT:\n",
      "************************************\n",
      " ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( (\n",
      " (\n",
      " ( (\n",
      "\n",
      "\n",
      " F F 35 (\n",
      " F\n",
      "\n",
      "\n",
      " (\n",
      "\n",
      "\n",
      "\n",
      " 35\n",
      " 35 F ( F<|endoftext|> 35 F F F (\n",
      " (\n",
      " ( ( F ( 35 F ( F F\n",
      " 35 F (<|endoftext|> 35 F F 35 A\n",
      " 35 F A<|endoftext|>\n",
      " F<|endoftext|> F 35<|endoftext|><|endoftext|> F F (\n",
      "\n",
      " F\n",
      " F F 35 F 35 ( The\n",
      " F 35 F F F 35 I F 35 F 35 F F F A F ( F F ( F F F The F F F A F F F F A 35 F F The The F<|endoftext|> F A F F 35<|endoftext|> A<|endoftext|> F F I A F The F F The F The The F The I<|endoftext|> The F A The F F The F\n",
      " I The F F The The F 35 The F F<|endoftext|> 35 The\n",
      " The F The F The H The F The H 35 F The\n",
      " 35 The The A The The F\n",
      " 35 The<|endoftext|> F The F The<|endoftext|>\n",
      "************************************\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "prompt_str = \"Compared to spoons, forks are meant to:\"\n",
    "prompt_tokens = tokenizer.encode(prompt_str)\n",
    "print(\"Generating with TT:\")\n",
    "generate_with_tt(tt_model, prompt_tokens.copy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f4b67723",
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt_str = \"Bees are similar to:\"\n",
    "prompt_tokens = tokenizer.encode(prompt_str)\n",
    "print(\"Generating with TT:\")\n",
    "generate_with_tt(tt_model, prompt_tokens.copy())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "efe3330e",
   "metadata": {},
   "source": [
    "Now try your own prompt!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d8beb647",
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt_str = input(\"Enter your prompt: \")\n",
    "\n",
    "prompt_tokens = tokenizer.encode(prompt_str)\n",
    "print(\"Generating with TT model:\")\n",
    "generate_with_tt(tt_model, prompt_tokens)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".tenstorrent-venv (3.10.12)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
